package com.example.shuiyin.aspect;

import com.example.shuiyin.exception.AppException;
import com.example.shuiyin.exception.ErrorCode;
import com.google.common.util.concurrent.RateLimiter;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import jakarta.servlet.http.HttpServletRequest;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

@Aspect
@Component
public class RateLimiterAspect {

    @Autowired
    private ConcurrentHashMap<String, RateLimiter> rateLimiterMap;

    @Pointcut("execution(* com.example.shuiyin.controller.*.*(..))")
    public void rateLimiterPointcut() {}

    @Around("rateLimiterPointcut()")
    public Object around(ProceedingJoinPoint point) throws Throwable {
        HttpServletRequest request = ((ServletRequestAttributes) RequestContextHolder.getRequestAttributes()).getRequest();
        String ip = getIpAddress(request);
        String uri = request.getRequestURI();
        String key = ip + ":" + uri;

        // 创建限流器或获取已有的
        RateLimiter limiter = rateLimiterMap.computeIfAbsent(key, k -> RateLimiter.create(5.0)); // 每秒5个请求

        // 尝试获取令牌，等待100ms
        if (!limiter.tryAcquire(1, 100, TimeUnit.MILLISECONDS)) {
            throw new AppException("请求过于频繁，请稍后再试", ErrorCode.TOO_MANY_REQUESTS);
        }

        return point.proceed();
    }

    private String getIpAddress(HttpServletRequest request) {
        String ip = request.getHeader("X-Forwarded-For");
        if (ip == null || ip.length() == 0 || "unknown".equalsIgnoreCase(ip)) {
            ip = request.getHeader("Proxy-Client-IP");
        }
        if (ip == null || ip.length() == 0 || "unknown".equalsIgnoreCase(ip)) {
            ip = request.getHeader("WL-Proxy-Client-IP");
        }
        if (ip == null || ip.length() == 0 || "unknown".equalsIgnoreCase(ip)) {
            ip = request.getRemoteAddr();
        }
        // 多代理情况下，第一个IP为客户端真实IP
        if (ip != null && ip.indexOf(",") > 0) {
            ip = ip.substring(0, ip.indexOf(","));
        }
        return ip;
    }
} 